#!/usr/bin/env python
# -*- coding:utf-8 -*-
import shelve
import subprocess
import os
import capstone as cs
import platform

PATH = os.path.dirname(os.path.realpath(__file__))
TEMP = os.path.join(PATH,'temp')
NASM = os.path.join(PATH,"nasm.exe")
TMPA = os.path.join(TEMP,"asmcode.asm")
TMPO = os.path.join(TEMP,"asmcode.out")


def get_local_arch():
    arch = platform.machine()
    return arch

def write_asm_to_file(asmcode):
    try:
        with open(TMPA,"w") as fwrite:
            fwrite.write(asmcode)
    except Exception as e:
        return False
    return True

def get_hex_code(asmcode):
    hexCode = b""
    if write_asm_to_file(asmcode):
        if subprocess.call([NASM,'-f','bin',TMPA,'-o',TMPO]) == 0:#subprocess.call([NASM,'-f','bin',TMPA,'-o',TMPO])
            hexCode = open(TMPO,"rb").read()
            #os.remove(TMPA)
            #os.remove(TMPO)
    return hexCode

def disasm(code,arch):
    address = 0x100000
    if arch == "I386":
        md = cs.Cs(cs.CS_ARCH_X86, cs.CS_MODE_32)
    elif arch == "AMD64":
        md = cs.Cs(cs.CS_ARCH_X86, cs.CS_MODE_64)
    elif arch == "MIPS":
        md = cs.Cs(cs.CS_ARCH_X86, cs.CS_MODE_32 + cs.CS_MODE_BIG_ENDIAN)
    elif arch == "ARM":
        md = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_ARM)
    else:
        md = cs.Cs(cs.CS_ARCH_X86, cs.CS_MODE_32)
    md.detail = True
    md.skipdata_setup = ("db", None, None)
    md.skipdata = True
    for i in md.disasm(code,address):
        print("0x%x:  %s  %s" %(i.address, i.mnemonic, i.op_str))
